#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试图数据库功能恢复
验证 find_related_concepts 和 graph_query 工具是否正常工作
"""

import asyncio
import sys
import os
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / "src"))

# 设置环境变量
os.environ['PYTHONPATH'] = str(project_root / "src")

async def test_graph_database_recovery():
    """
    测试图数据库功能恢复
    """
    print("=== 测试图数据库功能恢复 ===\n")
    
    try:
        # 1. 导入必要模块
        print("1. 导入模块...")
        from src.core.context_orchestrator import ContextOrchestrator
        from src.config import load_config
        from src.mcp.tools import MCPTools
        
        # 2. 初始化系统
        print("\n2. 初始化系统...")
        config = load_config()
        context_orchestrator = ContextOrchestrator(config)
        await context_orchestrator.initialize()
        
        neo4j_client = context_orchestrator.db_manager.neo4j_client
        print("✓ Neo4j 连接成功")
        
        # 初始化MCP工具
        mcp_tools = MCPTools(context_orchestrator)
        print("✓ MCP工具初始化成功")
        
        # 3. 测试基础数据查询
        print("\n3. 测试基础数据查询...")
        
        # 查询所有概念节点
        all_concepts_query = "MATCH (c:Concept) RETURN c.concept_id, c.name LIMIT 10"
        all_concepts = await neo4j_client.execute_query(all_concepts_query)
        print(f"✓ 数据库中共有 {len(all_concepts)} 个概念节点")
        
        if all_concepts:
            print("概念列表:")
            for record in all_concepts:
                print(f"  - ID: {record['c.concept_id']}, 名称: {record['c.name']}")
        else:
            print("❌ 数据库中没有概念节点")
            return False
        
        # 4. 测试 find_related_concepts 工具
        print("\n4. 测试 find_related_concepts 工具...")
        
        # 测试用例1: 查找 vue 相关概念
        print("\n测试用例1: 查找 vue 相关概念")
        try:
            vue_related_result = await mcp_tools._find_related_concepts({
                "concept": "vue"
            })
            if vue_related_result["success"]:
                vue_related = vue_related_result["results"]
                print(f"✓ 找到 {len(vue_related)} 个与 vue 相关的概念")
                for concept in vue_related:
                    print(f"  - {concept}")
            else:
                print(f"❌ find_related_concepts 返回错误: {vue_related_result.get('error')}")
                return False
        except Exception as e:
            print(f"❌ find_related_concepts 测试失败: {e}")
            return False
        
        # 测试用例2: 查找不存在的概念
        print("\n测试用例2: 查找不存在的概念")
        try:
            nonexistent_result = await mcp_tools._find_related_concepts({
                "concept": "nonexistent_concept"
            })
            if nonexistent_result["success"]:
                nonexistent_related = nonexistent_result["results"]
                print(f"✓ 查找不存在概念返回 {len(nonexistent_related)} 个结果（预期为0）")
            else:
                print(f"❌ find_related_concepts 返回错误: {nonexistent_result.get('error')}")
                return False
        except Exception as e:
            print(f"❌ find_related_concepts 测试失败: {e}")
            return False
        
        # 5. 测试 graph_query 工具
        print("\n5. 测试 graph_query 工具...")
        
        # 测试用例1: 查询所有概念
        print("\n测试用例1: 查询所有概念")
        try:
            all_concepts_query = "MATCH (c:Concept) RETURN c.concept_id, c.name LIMIT 5"
            graph_result_response = await mcp_tools._graph_query({
                "query": all_concepts_query
            })
            if graph_result_response["success"]:
                graph_result = graph_result_response["result"]
                print(f"✓ graph_query 返回 {len(graph_result)} 个结果")
                for result in graph_result:
                    print(f"  - {result}")
            else:
                print(f"❌ graph_query 返回错误: {graph_result_response.get('error')}")
                return False
        except Exception as e:
            print(f"❌ graph_query 测试失败: {e}")
            return False
        
        # 测试用例2: 查询概念关系
        print("\n测试用例2: 查询概念关系")
        try:
            relationship_query = """
            MATCH (c1:Concept)-[r]->(c2:Concept)
            RETURN c1.name, type(r), c2.name
            LIMIT 5
            """
            rel_result_response = await mcp_tools._graph_query({
                "query": relationship_query
            })
            if rel_result_response["success"]:
                rel_result = rel_result_response["result"]
                print(f"✓ 找到 {len(rel_result)} 个概念关系")
                for result in rel_result:
                    print(f"  - {result}")
            else:
                print(f"❌ graph_query 返回错误: {rel_result_response.get('error')}")
                return False
        except Exception as e:
            print(f"❌ graph_query 关系测试失败: {e}")
            return False
        
        # 6. 测试特定概念的详细信息
        print("\n6. 测试特定概念的详细信息...")
        
        try:
            vue_detail_query = """
            MATCH (c:Concept {concept_id: 'vue'})
            RETURN c.concept_id, c.name, c.description, c.category, c.tags
            """
            vue_detail_response = await mcp_tools._graph_query({
                "query": vue_detail_query
            })
            
            if not vue_detail_response["success"]:
                print(f"❌ graph_query 返回错误: {vue_detail_response.get('error')}")
                return False
            
            vue_detail = vue_detail_response["result"]
            
            if vue_detail:
                print("✓ Vue概念详细信息:")
                detail = vue_detail[0]
                print(f"  ID: {detail.get('c.concept_id')}")
                print(f"  名称: {detail.get('c.name')}")
                print(f"  描述: {detail.get('c.description')}")
                print(f"  分类: {detail.get('c.category')}")
                print(f"  标签: {detail.get('c.tags')}")
            else:
                print("❌ 未找到 Vue 概念详细信息")
                return False
        except Exception as e:
            print(f"❌ 概念详细信息测试失败: {e}")
            return False
        
        # 7. 测试概念搜索功能
        print("\n7. 测试概念搜索功能...")
        
        try:
            # 测试按名称搜索
            search_results = await neo4j_client.find_concepts_by_name("vue")
            print(f"✓ 按名称搜索 'vue' 找到 {len(search_results)} 个结果")
            
            # 测试查找相关概念
            related_concepts = await neo4j_client.find_related_concepts("vue")
            print(f"✓ 查找与 'vue' 相关的概念找到 {len(related_concepts)} 个结果")
            
        except Exception as e:
            print(f"❌ 概念搜索功能测试失败: {e}")
            return False
        
        # 8. 关闭连接
        print("\n8. 关闭连接...")
        await context_orchestrator.shutdown()
        print("✓ 连接已关闭")
        
        print("\n=== 所有测试通过！图数据库功能已恢复正常 ===\n")
        
        # 9. 输出测试总结
        print("测试总结:")
        print("✓ Neo4j 数据库连接正常")
        print("✓ 概念节点数据存在")
        print("✓ find_related_concepts 工具正常工作")
        print("✓ graph_query 工具正常工作")
        print("✓ 概念关系查询正常")
        print("✓ 概念搜索功能正常")
        print("\n现在可以正常使用 MCP 图数据库相关工具了！")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试过程中出错: {e}")
        import traceback
        traceback.print_exc()
        
        # 确保连接关闭
        try:
            await context_orchestrator.shutdown()
        except:
            pass
        
        return False

if __name__ == "__main__":
    # 运行测试
    success = asyncio.run(test_graph_database_recovery())
    
    if success:
        print("\n🎉 图数据库功能恢复测试成功！")
        sys.exit(0)
    else:
        print("\n❌ 图数据库功能恢复测试失败！")
        sys